// Copyright 2008 Isis Innovation Limited
#include "OpenGL.h"
#include <gvars3/instances.h>
#include "CameraCalibrator.h"
#include <TooN/SVD.h>
#include <cvd/colourspace_convert.h>
#include <fstream>
#include <stdlib.h>
#include <dirent.h>
#include <string>
#include <opencv/highgui.h>
#include <boost/lexical_cast.hpp>

using namespace CVD;
using namespace std;
using namespace GVars3;

/*
static void mavlink_handler(const lcm_recv_buf_t *rbuf, const char * channel, const mavlink_message_t* msg, void * user)
{
	switch(msg->msgid)
	{
	  case MAVLINK_MSG_ID_IMAGE_AVAILABLE:
	  {
	    CameraCalibrator *c = static_cast<CameraCalibrator*>(user);
      c->NewFrame(msg);
	  }
	  break;

	  default: //Ignore anything else
	  break;
	}
}
*/
void StringToDouble(const string str, double &number)
{
  stringstream buff;

  buff << str;
  buff >> number;
}


void DoubleToString(double number, string &str)
{
  stringstream buff;

  buff << number;
  buff >> str;
}

void GetImageList(string dir_name, string file_prefix, string file_type, vector<string>& file_list)
{
  DIR *dp;
  struct dirent *dirp;
  string filename;
  string extension;
  vector<double> frame_number;
  double number;
  unsigned int extension_pos, prefix_pos;
  int prefix_length = file_prefix.length();
	bool leading_zeros = false;
	int number_digits = 0;
  stringstream buff;

  if((dp  = opendir(dir_name.c_str())) == NULL) {
    cout << "Error: Unable to open directory " << dir_name << endl;
    exit(1);
  }

  while ((dirp = readdir(dp)) != NULL) {
    filename = string(dirp->d_name);
    extension_pos = filename.find(".");
    prefix_pos = filename.find(file_prefix);

    if(extension_pos != string::npos && prefix_pos != string::npos)
    {
      extension = filename.substr(extension_pos+1, string::npos);

			if(extension == file_type)
			{
		    filename.erase(extension_pos, string::npos); //erase extension
		    filename.erase(0, prefix_length); //erase 'Frame'

				if(filename.compare(0,1,"0") == 0)
				{
					leading_zeros = true;
					number_digits = filename.length();
				}
		    StringToDouble(filename, number);
		    frame_number.push_back(number);
			}
    }
  }
  closedir(dp);

  sort(frame_number.begin(), frame_number.end());

  for(vector<double>::iterator iter = frame_number.begin(); iter < frame_number.end(); iter++)
  {
    DoubleToString(*iter, filename);
		if(leading_zeros)
			filename.insert(0, number_digits-filename.length(), '0');
    filename.insert(0, file_prefix);
    filename.insert(0, dir_name);
    filename.append(".").append(extension);
    file_list.push_back(filename);
  }
}

int main(int argc, char **argv)
{
  ros::init(argc,argv,"uavdatmocalibrator");
  ros::NodeHandle n;
  ros::Subscriber ros_image_subscriber;
  cout << "  Welcome to CameraCalibrator " << endl;
  cout << "  -------------------------------------- " << endl;
  cout << "  Parallel tracking and mapping for Small AR workspaces" << endl;
  cout << "  Copyright (C) Isis Innovation Limited 2008 " << endl;
  cout << endl;  

  GUI.LoadFile("./config/settings.cfg");
  GUI.StartParserThread();
  atexit(GUI.StopParserThread); // Clean up readline when program quits
  
  gvar3<Vector<NUMTRACKERCAMPARAMETERS> > campam("Camera.Parameters");
  *campam = ATANCamera::mvDefaultParams; // Ignore camera parameters currently in file...
  static gvar3<int> gvnUseCamera("Calibrator.UseCamera", 1, SILENT); 

  if(*gvnUseCamera)
  {

    try
    {
      CameraCalibrator *c = new CameraCalibrator();
      ros_image_subscriber = n.subscribe("camera/rgb/image_color",1000,&CameraCalibrator::NewFrame,c);


      while(!c->Stop())
      {
        // Block waiting for new messages, when a message is received it will be processed
        ros::spin();
      }


      exit(EXIT_SUCCESS);
    }
    catch(CVD::Exceptions::All e)
    {
      cout << "ERROR: Failed to run CameraCalibrator, got exception: " << e.what << endl;
    }

  }
  else
  {
    CameraCalibrator c;
    static gvar3<string> gvsDirectory("Calibrator.ImageDirectory", "", SILENT);
    static gvar3<string> gvsPrefix("Calibrator.ImagePrefix", "", SILENT);
    static gvar3<string> gvsExtension("Calibrator.ImageExtension", "", SILENT);

    vector<string> images;
    GetImageList(*gvsDirectory, *gvsPrefix, *gvsExtension, images);

    for(int i = 0; i < images.size(); ++i)//i += 40
      c.NewImage(images[i]);

    while(!c.Stop())
      c.OptimizeImages();
  }
}


CameraCalibrator::CameraCalibrator() : mCamera("Camera")
{
  mbDone = false;

  GVars3::GV3::Register(mgvnCameraWidth, "Camera.Width", 640, GVars3::HIDDEN | GVars3::FATAL_IF_NOT_DEFINED);
  GVars3::GV3::Register(mgvnCameraHeight, "Camera.Height", 480, GVars3::HIDDEN | GVars3::FATAL_IF_NOT_DEFINED);
  GVars3::GV3::Register(mgvnCameraDepth, "Camera.Depth", 8, GVars3::HIDDEN | GVars3::FATAL_IF_NOT_DEFINED);
  GVars3::GV3::Register(mgvnCameraChannels, "Camera.Channels", 3, GVars3::HIDDEN | GVars3::FATAL_IF_NOT_DEFINED);
  GVars3::GV3::Register(mgvnCameraFrameRate, "Camera.FrameRate", 60, GVars3::HIDDEN | GVars3::FATAL_IF_NOT_DEFINED);

  mpGLWindow = new GLWindow2(ImageRef(*mgvnCameraWidth, *mgvnCameraHeight), "Camera Calibrator");

  mpIplImage = cvCreateImage(cvSize(*mgvnCameraWidth, *mgvnCameraHeight), *mgvnCameraDepth, *mgvnCameraChannels);
  mirFrameSize = ImageRef(*mgvnCameraWidth, *mgvnCameraHeight);
  mimFrameBW.resize(mirFrameSize);
  mimFrameRGB.resize(mirFrameSize);

  GUI.RegisterCommand("CameraCalibrator.GrabNextFrame", GUICommandCallBack, this);
  GUI.RegisterCommand("CameraCalibrator.Reset", GUICommandCallBack, this);
  GUI.RegisterCommand("CameraCalibrator.ShowNext", GUICommandCallBack, this);
  GUI.RegisterCommand("CameraCalibrator.SaveCalib", GUICommandCallBack, this);
  GUI.RegisterCommand("quit", GUICommandCallBack, this);
  GUI.RegisterCommand("exit", GUICommandCallBack, this);
  GV3::Register(mgvnOptimizing, "CameraCalibrator.Optimize", 0, SILENT);
  GV3::Register(mgvnShowImage, "CameraCalibrator.Show", 0, SILENT);
  GV3::Register(mgvnDisableDistortion, "CameraCalibrator.NoDistortion", 0, SILENT);

  GUI.ParseLine("GLWindow.AddMenu CalibMenu");
  GUI.ParseLine("CalibMenu.AddMenuButton Live GrabFrame CameraCalibrator.GrabNextFrame");
  GUI.ParseLine("CalibMenu.AddMenuButton Live Reset CameraCalibrator.Reset");
  GUI.ParseLine("CalibMenu.AddMenuButton Live Optimize \"CameraCalibrator.Optimize=1\"");
  GUI.ParseLine("CalibMenu.AddMenuToggle Live NoDist CameraCalibrator.NoDistortion");
  GUI.ParseLine("CalibMenu.AddMenuSlider Opti \"Show Img\" CameraCalibrator.Show 0 10");
  GUI.ParseLine("CalibMenu.AddMenuButton Opti \"Show Next\" CameraCalibrator.ShowNext");
  GUI.ParseLine("CalibMenu.AddMenuButton Opti \"Grab More\" CameraCalibrator.Optimize=0 ");
  GUI.ParseLine("CalibMenu.AddMenuButton Opti Reset CameraCalibrator.Reset");
  GUI.ParseLine("CalibMenu.AddMenuToggle Opti NoDist CameraCalibrator.NoDistortion");
  GUI.ParseLine("CalibMenu.AddMenuButton Opti Save CameraCalibrator.SaveCalib");

  Reset();
}

CameraCalibrator::~CameraCalibrator()
{
  cvReleaseImage(&mpIplImage);

}

bool CameraCalibrator::Stop()
{
  return mbDone;
}

void CameraCalibrator::NewImage(const string filename)
{
  IplImage *img = cvLoadImage(filename.c_str(), CV_LOAD_IMAGE_GRAYSCALE);

  // Copy from IplImage to CVD Image
  BasicImage<CVD::byte> bim((CVD::byte*)(img->imageData), mirFrameSize);
  mimFrameBW.copy_from(bim);

  // Set up openGL
  mpGLWindow->SetupViewport();
  mpGLWindow->SetupVideoOrtho();
  mpGLWindow->SetupVideoRasterPosAndZoom();
   
  GUI.ParseLine("CalibMenu.ShowMenu Live");
  glRasterPos2i(0,0);
  glDrawPixels(mimFrameBW.size().x, mimFrameBW.size().y, GL_LUMINANCE, GL_UNSIGNED_BYTE, mimFrameBW.data());
  
  CalibImage c;
  if(c.MakeFromImage(mimFrameBW))
  {
    mvCalibImgs.push_back(c);
    mvCalibImgs.back().GuessInitialPose(mCamera);
    mvCalibImgs.back().Draw3DGrid(mCamera, false);
  }
    
  ostringstream ost;
  ost << "Camera Calibration: Grabbed " << mvCalibImgs.size() << " images." << endl;
  ost << "Take snapshots of the calib grid with the \"GrabFrame\" button," << endl;
  ost << "and then press \"Optimize\"." << endl;
  ost << "Take enough shots (4+) at different angles to get points " << endl;
  ost << "into all parts of the image (corners too.) The whole grid " << endl;
  ost << "doesn't need to be visible so feel free to zoom in." << endl;

  mpGLWindow->DrawCaption(ost.str());
  mpGLWindow->DrawMenus();
  mpGLWindow->HandlePendingEvents();
  mpGLWindow->swap_buffers();
  cvReleaseImage(&img);

  //sleep(1);
}

void CameraCalibrator::OptimizeImages()
{
  // Set up openGL
  mpGLWindow->SetupViewport();
  mpGLWindow->SetupVideoOrtho();
  mpGLWindow->SetupVideoRasterPosAndZoom();

  OptimizeOneStep();
    
  GUI.ParseLine("CalibMenu.ShowMenu Opti");
  int nToShow = *mgvnShowImage - 1;
  if(nToShow < 0)
    nToShow = 0;
  if(nToShow >= (int) mvCalibImgs.size())
    nToShow = mvCalibImgs.size()-1;
  *mgvnShowImage = nToShow + 1;
    
  //glDrawPixels(mvCalibImgs[nToShow].mim);
  glRasterPos2i(0,0);
  glDrawPixels(mvCalibImgs[nToShow].mim.size().x, mvCalibImgs[nToShow].mim.size().y, GL_LUMINANCE, GL_UNSIGNED_BYTE, mvCalibImgs[nToShow].mim.data());

  mvCalibImgs[nToShow].Draw3DGrid(mCamera,true);
    
  ostringstream ost;
  ost << "Camera Calibration: Grabbed " << mvCalibImgs.size() << " images." << endl;
  ost << "Current RMS pixel error is " << mdMeanPixelError << endl;
  ost << "Current camera params are  " << GV3::get_var("Camera.Parameters") << endl;
  ost << "(That would be a pixel aspect ratio of " 
      <<  mCamera.PixelAspectRatio() << ")" << endl;
  ost << "Check fit by looking through the grabbed images." << endl;
  ost << "RMS should go below 0.5, typically below 0.3 for a wide lens." << endl;
  ost << "Press \"save\" to save calibration to camera.cfg file and exit." << endl;

  mpGLWindow->DrawCaption(ost.str());
  mpGLWindow->DrawMenus();
  mpGLWindow->HandlePendingEvents();
  mpGLWindow->swap_buffers();
}

void CameraCalibrator::NewFrame(const sensor_msgs::Image::ConstPtr &msg)
{

  // Copy from IplImage to CVD Image
  if(*mgvnCameraChannels == 1)
  {
    memcpy(mpIplImage->imageDataOrigin,msg->data.data(),msg->height*msg->width);
    BasicImage<CVD::byte> bim((CVD::byte*)(mpIplImage->imageData), mirFrameSize);
    mimFrameBW.copy_from(bim);
  }
  else
  {
    memcpy(mpIplImage->imageDataOrigin,msg->data.data(),msg->height*msg->width*3);
   // cvConvertImage(mpIplImage, mpIplImage, CV_CVTIMG_SWAP_RB); //Opencv data is in BGR order, swap R & B to get RGB order
    BasicImage<CVD::Rgb<CVD::byte> > bim((CVD::Rgb<CVD::byte>*)(mpIplImage->imageData), mirFrameSize);
    mimFrameRGB.copy_from(bim);
    mimFrameBW = convert_image(mimFrameRGB);
  }

  // Set up openGL
  mpGLWindow->SetupViewport();
  mpGLWindow->SetupVideoOrtho();
  mpGLWindow->SetupVideoRasterPosAndZoom();
    
  if(mvCalibImgs.size() < 1)
    *mgvnOptimizing = 0;
    
  if(!*mgvnOptimizing)
  {
    GUI.ParseLine("CalibMenu.ShowMenu Live");
    //glDrawPixels(mimFrameBW);
    glRasterPos2i(0,0);
    glDrawPixels(mimFrameBW.size().x, mimFrameBW.size().y, GL_LUMINANCE, GL_UNSIGNED_BYTE, mimFrameBW.data());
  
    CalibImage c;
    if(c.MakeFromImage(mimFrameBW))
    {
      if(mbGrabNextFrame)
	    {
	      mvCalibImgs.push_back(c);
	      mvCalibImgs.back().GuessInitialPose(mCamera);
	      mvCalibImgs.back().Draw3DGrid(mCamera, false);
	      mbGrabNextFrame = false;
	    };
    }
  }
  else
  {
    OptimizeOneStep();
      
    GUI.ParseLine("CalibMenu.ShowMenu Opti");
    int nToShow = *mgvnShowImage - 1;
    if(nToShow < 0)
      nToShow = 0;
    if(nToShow >= (int) mvCalibImgs.size())
      nToShow = mvCalibImgs.size()-1;
    *mgvnShowImage = nToShow + 1;
      
    //glDrawPixels(mvCalibImgs[nToShow].mim);
    glRasterPos2i(0,0);
    glDrawPixels(mvCalibImgs[nToShow].mim.size().x, mvCalibImgs[nToShow].mim.size().y, GL_LUMINANCE, GL_UNSIGNED_BYTE, mvCalibImgs[nToShow].mim.data());
    mvCalibImgs[nToShow].Draw3DGrid(mCamera,true);
  }
    
  ostringstream ost;
  ost << "Camera Calibration: Grabbed " << mvCalibImgs.size() << " images." << endl;
  if(!*mgvnOptimizing)
  {
    ost << "Take snapshots of the calib grid with the \"GrabFrame\" button," << endl;
    ost << "and then press \"Optimize\"." << endl;
    ost << "Take enough shots (4+) at different angles to get points " << endl;
    ost << "into all parts of the image (corners too.) The whole grid " << endl;
    ost << "doesn't need to be visible so feel free to zoom in." << endl;
  }
  else
  {
    ost << "Current RMS pixel error is " << mdMeanPixelError << endl;
    ost << "Current camera params are  " << GV3::get_var("Camera.Parameters") << endl;
    ost << "(That would be a pixel aspect ratio of " 
        <<  mCamera.PixelAspectRatio() << ")" << endl;
    ost << "Check fit by looking through the grabbed images." << endl;
    ost << "RMS should go below 0.5, typically below 0.3 for a wide lens." << endl;
    ost << "Press \"save\" to save calibration to camera.cfg file and exit." << endl;
  }

  mpGLWindow->DrawCaption(ost.str());
  mpGLWindow->DrawMenus();
  mpGLWindow->HandlePendingEvents();
  mpGLWindow->swap_buffers();
}

void CameraCalibrator::Reset()
{
  *mCamera.mgvvCameraParams = ATANCamera::mvDefaultParams;
  if(*mgvnDisableDistortion) mCamera.DisableRadialDistortion();
  mbGrabNextFrame =false;
  *mgvnOptimizing = false;
  mvCalibImgs.clear();
}

void CameraCalibrator::GUICommandCallBack(void* ptr, string sCommand, string sParams)
{
  ((CameraCalibrator*) ptr)->GUICommandHandler(sCommand, sParams);
}

void CameraCalibrator::GUICommandHandler(string sCommand, string sParams)  // Called by the callback func..
{
  if(sCommand=="CameraCalibrator.Reset")
    {
      Reset();
      return;
    };
  if(sCommand=="CameraCalibrator.GrabNextFrame")
    {
      mbGrabNextFrame = true;
      return;
    }
  if(sCommand=="CameraCalibrator.ShowNext")
    {
      int nToShow = (*mgvnShowImage - 1 + 1) % mvCalibImgs.size();
      *mgvnShowImage = nToShow + 1;
      return;
    }
  if(sCommand=="CameraCalibrator.SaveCalib")
    {
      cout << "  Camera calib is " << GV3::get_var("Camera.Parameters") << endl;
      cout << "  Saving camera calib to camera.cfg..." << endl;
      ofstream ofs("camera.cfg");
      if(ofs.good())
	{
	  GV2.PrintVar("Camera.Parameters", ofs);
	  GV2.PrintVar("Camera.Width", ofs);
	  GV2.PrintVar("Camera.Height", ofs);
	  GV2.PrintVar("Camera.Depth", ofs);
	  GV2.PrintVar("Camera.Channels", ofs);
	  GV2.PrintVar("Camera.FrameRate", ofs);
	  ofs.close();

    IplImage* img = cvCreateImageHeader(cvSize(mvCalibImgs[0].mim.size().x, mvCalibImgs[0].mim.size().y), 8, 1);    
    for(int i = 0; i < mvCalibImgs.size(); ++i)
    {
      cvSetData(img, mvCalibImgs[i].mim.data(), mvCalibImgs[i].mim.row_stride());
      string filename = "calibration_" + boost::lexical_cast<std::string>(i) + ".jpg";
      cvSaveImage(filename.c_str(), img);
    }
    cvReleaseImageHeader(&img);

	  cout << "  .. saved."<< endl;
	}
      else
	{
	  cout <<"! Could not open camera.cfg for writing." << endl;
	  GV2.PrintVar("Camera.Parameters", cout);
	  GV2.PrintVar("Camera.Width", cout);
	  GV2.PrintVar("Camera.Height", cout);
	  GV2.PrintVar("Camera.Depth", cout);
	  GV2.PrintVar("Camera.Channels", cout);
	  GV2.PrintVar("Camera.FrameRate", cout);
	  cout <<"  Copy-paste above lines to settings.cfg or camera.cfg! " << endl;
	}
      mbDone = true;
    }
  if(sCommand=="exit" || sCommand=="quit")
    {
      mbDone = true;
    }
}

void CameraCalibrator::OptimizeOneStep()
{
  int nViews = mvCalibImgs.size();
  int nDim = 6 * nViews + NUMTRACKERCAMPARAMETERS;
  int nCamParamBase = nDim - NUMTRACKERCAMPARAMETERS;
  
  Matrix<> mJTJ(nDim, nDim);
  Vector<> vJTe(nDim);
  mJTJ = Identity; // Weak stabilizing prior
  vJTe = Zeros;

  if(*mgvnDisableDistortion) mCamera.DisableRadialDistortion();

  
  double dSumSquaredError = 0.0;
  int nTotalMeas = 0;
  
  for(int n=0; n<nViews; n++)
    {
      int nMotionBase = n*6;
      vector<CalibImage::ErrorAndJacobians> vEAJ = mvCalibImgs[n].Project(mCamera);
      for(unsigned int i=0; i<vEAJ.size(); i++)
	{
	  CalibImage::ErrorAndJacobians &EAJ = vEAJ[i];
	  // All the below should be +=, but the MSVC compiler doesn't seem to understand that. :(
      mJTJ.slice(nMotionBase, nMotionBase, 6, 6) = 
      mJTJ.slice(nMotionBase, nMotionBase, 6, 6) + EAJ.m26PoseJac.T() * EAJ.m26PoseJac;
      mJTJ.slice(nCamParamBase, nCamParamBase, NUMTRACKERCAMPARAMETERS, NUMTRACKERCAMPARAMETERS) = 
      mJTJ.slice(nCamParamBase, nCamParamBase, NUMTRACKERCAMPARAMETERS, NUMTRACKERCAMPARAMETERS) + EAJ.m2NCameraJac.T() * EAJ.m2NCameraJac;
      mJTJ.slice(nMotionBase, nCamParamBase, 6, NUMTRACKERCAMPARAMETERS) =
      mJTJ.slice(nMotionBase, nCamParamBase, 6, NUMTRACKERCAMPARAMETERS) + EAJ.m26PoseJac.T() * EAJ.m2NCameraJac;
      mJTJ.T().slice(nMotionBase, nCamParamBase, 6, NUMTRACKERCAMPARAMETERS) = 
      mJTJ.T().slice(nMotionBase, nCamParamBase, 6, NUMTRACKERCAMPARAMETERS) + EAJ.m26PoseJac.T() * EAJ.m2NCameraJac;
      // Above does twice the work it needs to, but who cares..

      vJTe.slice(nMotionBase,6) = 
      vJTe.slice(nMotionBase,6) + EAJ.m26PoseJac.T() * EAJ.v2Error;
      vJTe.slice(nCamParamBase,NUMTRACKERCAMPARAMETERS) = 
      vJTe.slice(nCamParamBase,NUMTRACKERCAMPARAMETERS) + EAJ.m2NCameraJac.T() * EAJ.v2Error;

	  dSumSquaredError += EAJ.v2Error * EAJ.v2Error;
	  ++nTotalMeas;
	}
    };
  
  mdMeanPixelError = sqrt(dSumSquaredError / nTotalMeas);

  SVD<> svd(mJTJ);
  Vector<> vUpdate(nDim);
  vUpdate= svd.backsub(vJTe);
  vUpdate *= 0.1; // Slow down because highly nonlinear...
  for(int n=0; n<nViews; n++)
    mvCalibImgs[n].mse3CamFromWorld = SE3<>::exp(vUpdate.slice(n * 6, 6)) * mvCalibImgs[n].mse3CamFromWorld;
  mCamera.UpdateParams(vUpdate.slice(nCamParamBase, NUMTRACKERCAMPARAMETERS));
};















